import numpy as np

labels = {
        "unlabeled":0,
        "outlier":1,
        "car":10,
        "bicycle":11,
        "bus":13,
        "motorcycle":15,
        "on-rails":16,
        "truck":18,
        "other-vehicle":20,
        "person":30,
        "bicyclist":31,
        "motorcyclist":32,
        "road":40,
        "parking":44,
        "sidewalk":48,
        "other-ground":49,
        "building":50,
        "fence":51,
        "other-structure":52,
        "lane-marking":60,
        "vegetation":70,
        "trunk":71,
        "terrain":72,
        "pole":80,
        "traffic-sign":81,
        "other-object":99,
        "moving-car":252,
        "moving-bicyclist":253,
        "moving-person":254,
        "moving-motorcyclist":255,
        "moving-on-rails":256,
        "moving-bus":257,
        "moving-truck":258,
        "moving-other-vehicle":259,
        }

label2kimera_labels = {
        0:0,
        1:0,
        10:10,
        11:15,
        13:10,
        15:15,
        16:10,
        18:10,
        20:15,
        30:30,
        31:30,
        32:30,
        40:40,
        44:40,
        48:40,
        49:40,
        50:50,
        51:51,
        52:50,
        60:51,
        70:70,
        71:70,
        72:70,
        80:51,
        81:51,
        99:51,
        252:10,
        253:30,
        254:30,
        255:30,
        256:10,
        257:10,
        258:10,
        259:10,
}


colormap = {
        "unlabeled":np.array([255, 255, 255], dtype="uint8"),
        "outlier":np.array([0, 0, 0], dtype="uint8"),
        "car":np.array([0, 0, 255], dtype="uint8"),
        "bicycle":np.array([0, 255, 255], dtype="uint8"),
        "bus":np.array([0, 0, 255], dtype="uint8"),
        "motorcycle":np.array([0, 255, 255], dtype="uint8"),
        "on-rails":np.array([0, 0, 255], dtype="uint8"),
        "truck":np.array([0, 0, 255], dtype="uint8"),
        "other-vehicle":np.array([0, 255, 255], dtype="uint8"),
        "person":np.array([127, 0, 255], dtype="uint8"),
        "bicyclist":np.array([127, 0, 255], dtype="uint8"),
        "motorcyclist":np.array([127, 0, 255], dtype="uint8"),
        "road":np.array([255, 0, 127], dtype="uint8"),
        "parking":np.array([255, 0, 127], dtype="uint8"),
        "sidewalk":np.array([255, 0, 127], dtype="uint8"),
        "other-ground":np.array([255, 0, 127], dtype="uint8"),
        "building":np.array([255, 255, 0], dtype="uint8"),
        "fence":np.array([255, 255, 0], dtype="uint8"),
        "other-structure":np.array([255, 255, 0], dtype="uint8"),
        "lane-marking":np.array([255, 127, 0], dtype="uint8"),
        "vegetation":np.array([0, 255, 0], dtype="uint8"),
        "trunk":np.array([0, 255, 0], dtype="uint8"),
        "terrain":np.array([0, 255, 0], dtype="uint8"),
        "pole":np.array([255, 255, 0], dtype="uint8"),
        "traffic-sign":np.array([255, 255, 0], dtype="uint8"),
        "other-object":np.array([255, 255, 0], dtype="uint8"),
        "moving-car":np.array([0, 0, 0], dtype="uint8"),
        "moving-bicyclist":np.array([0, 0, 0], dtype="uint8"),
        "moving-person":np.array([0, 0, 0], dtype="uint8"),
        "moving-motorcyclist":np.array([0, 0, 0], dtype="uint8"),
        "moving-on-rails":np.array([0, 0, 0], dtype="uint8"),
        "moving-bus":np.array([0, 0, 0], dtype="uint8"),
        "moving-truck":np.array([0, 0, 0], dtype="uint8"),
        "moving-other-vehicle":np.array([0, 0, 0], dtype="uint8"),
}

colormap_a = {
        "unlabeled":np.array([255, 255, 255, 255], dtype="uint8"),
        "outlier":np.array([0, 0, 0, 255], dtype="uint8"),
        "car":np.array([0, 0, 255, 255], dtype="uint8"),
        "bicycle":np.array([0, 255, 255, 255], dtype="uint8"),
        "bus":np.array([0, 0, 255, 255], dtype="uint8"),
        "motorcycle":np.array([0, 255, 255, 255], dtype="uint8"),
        "on-rails":np.array([0, 0, 255, 255], dtype="uint8"),
        "truck":np.array([0, 0, 255, 255], dtype="uint8"),
        "other-vehicle":np.array([0, 255, 255, 255], dtype="uint8"),
        "person":np.array([127, 0, 255, 255], dtype="uint8"),
        "bicyclist":np.array([127, 0, 255, 255], dtype="uint8"),
        "motorcyclist":np.array([127, 0, 255, 255], dtype="uint8"),
        "road":np.array([255, 0, 127, 255], dtype="uint8"),
        "parking":np.array([255, 0, 127, 255], dtype="uint8"),
        "sidewalk":np.array([255, 0, 127, 255], dtype="uint8"),
        "other-ground":np.array([255, 0, 127, 255], dtype="uint8"),
        "building":np.array([255, 255, 0, 255], dtype="uint8"),
        "fence":np.array([255, 255, 0, 255], dtype="uint8"),
        "other-structure":np.array([255, 255, 0, 255], dtype="uint8"),
        "lane-marking":np.array([255, 127, 0, 255], dtype="uint8"),
        "vegetation":np.array([0, 255, 0, 255], dtype="uint8"),
        "trunk":np.array([0, 255, 0, 255], dtype="uint8"),
        "terrain":np.array([0, 255, 0, 255], dtype="uint8"),
        "pole":np.array([255, 255, 0, 255], dtype="uint8"),
        "traffic-sign":np.array([255, 255, 0, 255], dtype="uint8"),
        "other-object":np.array([255, 255, 0, 255], dtype="uint8"),
        "moving-car":np.array([0, 0, 0, 255], dtype="uint8"),
        "moving-bicyclist":np.array([0, 0, 0, 255], dtype="uint8"),
        "moving-person":np.array([0, 0, 0, 255], dtype="uint8"),
        "moving-motorcyclist":np.array([0, 0, 0, 255], dtype="uint8"),
        "moving-on-rails":np.array([0, 0, 0, 255], dtype="uint8"),
        "moving-bus":np.array([0, 0, 0, 255], dtype="uint8"),
        "moving-truck":np.array([0, 0, 0, 255], dtype="uint8"),
        "moving-other-vehicle":np.array([0, 0, 0, 255], dtype="uint8"),
}

def label2color(label_id):
    for k, v in labels.items():
        if v == label_id:
            return colormap[k]
    return np.array([0, 0, 0], dtype="uint8")

def label2color_a(label_id):
    for k, v in labels.items():
        if v == label_id:
            return colormap_a[k]
    return np.array([0, 0, 0], dtype="uint8")

get_label_id = np.vectorize(lambda l: int(l) & 0xFFFF)
